#include <btl/ViewGeometry/FivePointPose.hpp>
#include <btl/Maths/GaussJordanElimination.hpp>

#include <Eigen/Dense>
#include <Eigen/SVD>
#include <Eigen/Eigenvalues>

namespace btl
{
namespace viewgeometry
{

void FivePointStewenius::computeInternal(int opts)
{
   // Any linear combination of the 4 columns of nullspace will give you a matrix
   // which will meet the epipolar constraint for the points you're solving for.
   // However, there are only some points within this 4D space that represent
   // valid essential matrices.
   // This 4D space can be parameterised with 4 values (x,y,z,w), but since scale
   // is unobservable we fix w=1 and parameterise with just 3 values (x,y,z).
   //
   // We have some fixed constraints that come from the definition of an essential
   // matrix, specifically:
   //   (1)   det E = 0
   //   (2)   2EE^T E - trace(EE^T)E = [0]    (nb: [0] here is a 3x3 matrix of zeros)
   //
   // As described above, E is itself constructed from a linear combination of 4
   // basis E matrices, so we actually have:
   //   (1')  det (xE_1 + yE_2 + zE_3 + E_4) = 0
   //
   // and similarly for eq. (2).
   // This constraint can be re-written in terms of the individual coefficients of the
   // E_1, E_2, E_3 and E_4 matrices, and the three parameters x,y,z, to give a
   // polynomial in x,y,z.
   //   // For the the second constraint, we re-write as 9 separate constraints
   // (9 coefficients, each determined as a polynomial of x,y,z and each required to
   // evaluate to zero).
   //
   // There are 20 monomials in the polynomial constraint equations, corresponding
   // with all the different combinations of x, y and z parameters:
   // [x^3,x^2y,x^2z,xy^2,xyz,xz^2,y^3,y^2z,yz^2,z^3,x^2,xy,xz,y^2,yz,z^2,x,y,z,1]
   //
   // So, we have 10 constraints, in 20 monomials: these form the big M matrix.
   // Then we can solve for x, y and z, and use those values to reconstruct the
   // essential matrix (except actually we'll get 10 possible x,y,z combinations,
   // so we get 10 essential matrices, and we then have to do some extra checks
   // against more data to find the correct one).

   Eigen::JacobiSVD<Eigen::MatrixXd> svd(_epipolarConstraintMatrix.transpose(), Eigen::ComputeFullV);
   const Eigen::MatrixXd nullspace = svd.matrixV().block<9,4>(0,5);

   //std::cout << "Singular values: " << svd.singularValues().transpose() << "\n";
   //std::cout << "Null space:\n" << nullspace << "\n";

   const Eigen::MatrixXd M = buildMonomialMatrix(nullspace);

   //std::cout << "M:\n" << M.transpose() << "\n";

   btl::GaussJordanElimination<Eigen::MatrixXd> gaussJordan(M);
   const Eigen::MatrixXd B = gaussJordan.matrix().block<10,10>(0,10);

   //std::cout << "B':\n" << B.transpose() << "\n";

   // I have no idea why we use this particular combination of rows from the B matrix...
   // or why we set this particular set of elements to 1 in the rest of the At matrix...
   // I'll work it out eventually
   Eigen::MatrixXd At(10,10);
   At.row(0) = -B.row(0);
   At.row(1) = -B.row(1);
   At.row(2) = -B.row(2);
   At.row(3) = -B.row(4);
   At.row(4) = -B.row(5);
   At.row(5) = -B.row(7);
   At.row(6).setZero(); At(6,0) = 1.0;
   At.row(7).setZero(); At(7,1) = 1.0;
   At.row(8).setZero(); At(8,3) = 1.0;
   At.row(9).setZero(); At(9,6) = 1.0;

   //std::cout << "At':\n" << At.transpose() << "\n";

   Eigen::EigenSolver<Eigen::MatrixXd> eig(At);

   // extract the E null-space coordinates of the 10 (possibly imaginary) solutions
   Eigen::MatrixXcd sols(4,10);
   sols.row(0) = eig.eigenvectors().row(6).array() / eig.eigenvectors().row(9).array();
   sols.row(1) = eig.eigenvectors().row(7).array() / eig.eigenvectors().row(9).array();
   sols.row(2) = eig.eigenvectors().row(8).array() / eig.eigenvectors().row(9).array();
   sols.row(3).setOnes();

   Eigen::MatrixXcd Evec = nullspace * sols;
   assert(Evec.rows() == 9 && Evec.cols() == 10);

   _numSolutions = 0;

   // reconstruct (up to 10) possible E matrices
   for (int i = 0; i < 10; ++i)
   {
      bool has_imag = (Evec.col(i).array().imag() != 0.0).any();
      if (! has_imag)
      {
         Eigen::Matrix3d& sol = _solutions[_numSolutions++];

         Eigen::Matrix<double,9,1> Ev = Evec.real().col(i);
         sol.col(0) = Ev.middleRows(0,3);
         sol.col(1) = Ev.middleRows(3,3);
         sol.col(2) = Ev.middleRows(6,3);

         // rank fixup applies a loose check on the singular values
         // and then slightly adjusts the solution to have singular
         // values of exactly (alpha, alpha, 0) where alpha is the
         // average of the two singular values already found
         if ((opts & ApplyRankFix) == ApplyRankFix)
         {
            Eigen::JacobiSVD<Eigen::Matrix3d> svd(sol, Eigen::ComputeFullU | Eigen::ComputeFullV);
            Eigen::Vector3d sv = svd.singularValues();

            double svx;
            if ((opts & NormaliseResults) == NormaliseResults)
               svx = 1.0;
            else
               svx = (sv(0) + sv(1)) / 2.0;

            sol
               = svd.matrixU()
               * Eigen::Vector3d(svx, svx, 0.0).asDiagonal()
               * svd.matrixV().transpose();
         }
         else
         {
            if ((opts & NormaliseResults) == NormaliseResults)
            {
               double norm = std::sqrt(sol.array().abs2().rowwise().sum().sum());
               sol /= norm;
            }
         }
      }
   }
}

/// This performs the multiplication:
///   m = (a_x + a_y + a_z + a_1) * (b_x + b_y + b_z + b_1)
/// factoring the result by the combinations of parameters.
///
/// This is basically: a * b^T (which gives a 4x4 matrix), but with
/// equivalent combinations (e.g., a_x*b_y; a_y*b_x) combined
/// (which reduces it from 16 to 10 values)
Eigen::VectorXd mul2(const Eigen::Vector4d& a, const Eigen::Vector4d& b)
{
   Eigen::VectorXd m(10,1);
   m(0) = a(0)*b(0); // xx
   m(1) = a(0)*b(1) + a(1)*b(0); // xy, yx
   m(2) = a(0)*b(2) + a(2)*b(0); // xz, zx
   m(3) = a(1)*b(1); // yy
   m(4) = a(1)*b(2) + a(2)*b(1); // yz, zy
   m(5) = a(2)*b(2); // zz
   m(6) = a(0)*b(3) + a(3)*b(0); // x1, 1x
   m(7) = a(1)*b(3) + a(3)*b(1); // y1, 1y
   m(8) = a(2)*b(3) + a(3)*b(2); // z1, 1z
   m(9) = a(3)*b(3); // 11
   return m;
}

/// This performs the multiplication:
///   m = (a*b) * (c_x + c_y + c_z + c_1)
/// factoring the result by combinations of parameters, given that
/// (a*b) has already been produced by mul1
Eigen::VectorXd mul3(const Eigen::VectorXd& ab, const Eigen::Vector4d& c)
{
   Eigen::VectorXd m(20,1);
   m( 0) = ab(0)*c(0); // xxx
   m( 1) = ab(0)*c(1) + ab(1)*c(0); // xxy, xyx
   m( 2) = ab(1)*c(1) + ab(3)*c(0); // xyy, yyx
   m( 3) = ab(3)*c(1); // yyy
   m( 4) = ab(0)*c(2) + ab(2)*c(0); // xxz, xzx
   m( 5) = ab(1)*c(2) + ab(2)*c(1) + ab(4)*c(0); // xyz, xzy, yzx
   m( 6) = ab(3)*c(2) + ab(4)*c(1); // yyz, yzy
   m( 7) = ab(2)*c(2) + ab(5)*c(0); // xzz, zzx
   m( 8) = ab(4)*c(2) + ab(5)*c(1); // yzz, zzy
   m( 9) = ab(5)*c(2); // zzz
   m(10) = ab(0)*c(3) + ab(6)*c(0); // xx1, x1x
   m(11) = ab(1)*c(3) + ab(6)*c(1) + ab(7)*c(0); // xy1, x1y, y1x
   m(13) = ab(2)*c(3) + ab(6)*c(2) + ab(8)*c(0); // xz1, x1z, z1x
   m(12) = ab(3)*c(3) + ab(7)*c(1); // yy1, y1y
   m(14) = ab(4)*c(3) + ab(7)*c(2) + ab(8)*c(1); // yz1, y1z, z1y
   m(15) = ab(5)*c(3) + ab(8)*c(2); // zz1, z1z
   m(16) = ab(6)*c(3) + ab(9)*c(0); // x11, 11x
   m(17) = ab(7)*c(3) + ab(9)*c(1); // y11, 11y
   m(18) = ab(8)*c(3) + ab(9)*c(2); // z11, 11z
   m(19) = ab(9)*c(3); // 111
   return m;
}

Eigen::MatrixXd FivePointStewenius::buildMonomialMatrix(const Eigen::MatrixXd& Es)
{
   assert(Es.rows() == 9 && Es.cols() == 4);

   const Eigen::Vector4d e11 = Es.row(0).transpose();
   const Eigen::Vector4d e12 = Es.row(1).transpose();
   const Eigen::Vector4d e13 = Es.row(2).transpose();
   const Eigen::Vector4d e21 = Es.row(3).transpose();
   const Eigen::Vector4d e22 = Es.row(4).transpose();
   const Eigen::Vector4d e23 = Es.row(5).transpose();
   const Eigen::Vector4d e31 = Es.row(6).transpose();
   const Eigen::Vector4d e32 = Es.row(7).transpose();
   const Eigen::Vector4d e33 = Es.row(8).transpose();

   // build the trace constraint:
   //   2EE^T E - trace(EE^T)E
   // = 2EE^T E - trace(EE^T)I E
   // = (2EE^T - trace(EE^T)I) E
   // = (EE^T - 0.5 trace(EE^T)I) E

   const Eigen::VectorXd  EEt11 = mul2(e11,e11) + mul2(e12,e12) + mul2(e13,e13);
   const Eigen::VectorXd  EEt12 = mul2(e11,e21) + mul2(e12,e22) + mul2(e13,e23);
   const Eigen::VectorXd  EEt13 = mul2(e11,e31) + mul2(e12,e32) + mul2(e13,e33);
   const Eigen::VectorXd& EEt21 = EEt12;
   const Eigen::VectorXd  EEt22 = mul2(e21,e21) + mul2(e22,e22) + mul2(e23,e23);
   const Eigen::VectorXd  EEt23 = mul2(e21,e31) + mul2(e22,e32) + mul2(e23,e33);
   const Eigen::VectorXd& EEt31 = EEt13;
   const Eigen::VectorXd& EEt32 = EEt23;
   const Eigen::VectorXd  EEt33 = mul2(e31,e31) + mul2(e32,e32) + mul2(e33,e33);

   // Etr11 = (EE^T - 0.5 trace(EE^T)I)
   // (I = 3x3 identity matrix; therefore trace(EE^T)I is just trace(EE^T) in the diagonal elements)
   const Eigen::VectorXd traceEEt = 0.5 * (EEt11 + EEt22 + EEt33);
   const Eigen::VectorXd  Etr11 = EEt11 - traceEEt;
   const Eigen::VectorXd& Etr12 = EEt12;
   const Eigen::VectorXd& Etr13 = EEt13;
   const Eigen::VectorXd& Etr21 = EEt21;
   const Eigen::VectorXd  Etr22 = EEt22 - traceEEt;
   const Eigen::VectorXd& Etr23 = EEt23;
   const Eigen::VectorXd& Etr31 = EEt31;
   const Eigen::VectorXd& Etr32 = EEt32;
   const Eigen::VectorXd  Etr33 = EEt33 - traceEEt;

   Eigen::MatrixXd M(10,20);

   // build the final values of the trace constraint
   M.row(0) = mul3(Etr11,e11) + mul3(Etr12,e21) + mul3(Etr13,e31);
   M.row(1) = mul3(Etr21,e11) + mul3(Etr22,e21) + mul3(Etr23,e31);
   M.row(2) = mul3(Etr31,e11) + mul3(Etr32,e21) + mul3(Etr33,e31);
   M.row(3) = mul3(Etr11,e12) + mul3(Etr12,e22) + mul3(Etr13,e32);
   M.row(4) = mul3(Etr21,e12) + mul3(Etr22,e22) + mul3(Etr23,e32);
   M.row(5) = mul3(Etr31,e12) + mul3(Etr32,e22) + mul3(Etr33,e32);
   M.row(6) = mul3(Etr11,e13) + mul3(Etr12,e23) + mul3(Etr13,e33);
   M.row(7) = mul3(Etr21,e13) + mul3(Etr22,e23) + mul3(Etr23,e33);
   M.row(8) = mul3(Etr31,e13) + mul3(Etr32,e23) + mul3(Etr33,e33);

   // build the determinant constraint: det(E) = 0
   M.row(9)
      = mul3(mul2(e12,e23) - mul2(e13,e22), e31)
      + mul3(mul2(e13,e21) - mul2(e11,e23), e32)
      + mul3(mul2(e11,e22) - mul2(e12,e21), e33);

   return M;
}

} // namespace viewgeometry
} // namespace btl
